# 包含输入图png 100张
posenv_ori_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_ori' # control
# 包含预测图png 100张 + caption同名txt 100个
posenv_tar_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_tar' # folder
# 验证集输入 11张
posenv_val_ori_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_val_ori' # control
# 验证集gt 11张
posenv_val_tar_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_val_tar' # folder

# 输入-预测 文件夹 对应文件名相同

import os
import pandas as pd

# 配置路径
base_dir = "/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_processed"
os.makedirs(base_dir, exist_ok=True)
train_image_dir = os.path.join(base_dir, "images")
train_control_dir = os.path.join(base_dir, "control_images")
val_image_dir = os.path.join(base_dir, "val_images")
val_control_dir = os.path.join(base_dir, "val_control_images")
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_control_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(val_control_dir, exist_ok=True)

# 生成训练集 CSV 数据
train_data = []
for img_name in os.listdir(train_image_dir):
    if img_name.endswith(".png"):
        # 假设 caption 在同名 txt 文件中
        txt_path = os.path.join("/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_tar", img_name.replace(".png", ".txt"))
        caption = open(txt_path, "r").read().strip() if os.path.exists(txt_path) else ""
        train_data.append({
            "image": f"images/{img_name}",
            "control_image": f"control_images/{img_name}",
            "caption": caption
        })

# 生成验证集 CSV 数据
val_data = []
for img_name in os.listdir(val_image_dir):
    if img_name.endswith(".png"):
        txt_path = os.path.join("/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_val_tar", img_name.replace(".png", ".txt"))
        caption = open(txt_path, "r").read().strip() if os.path.exists(txt_path) else ""
        val_data.append({
            "image": f"val_images/{img_name}",
            "control_image": f"val_control_images/{img_name}",
            "caption": caption
        })



# 合并并保存 CSV
df = pd.DataFrame(train_data + val_data)
df.to_csv(os.path.join(base_dir, "metadata_kontext.csv"), index=False)
print("CSV 生成完成！")

import pandas as pd
import os

base_dir = "/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_processed"
csv_path = os.path.join(base_dir, "metadata_kontext.csv")
df = pd.read_csv(csv_path)

# 检查文件是否存在
for _, row in df.iterrows():
    assert os.path.exists(os.path.join(base_dir, row["image"])), f"Missing image: {row['image']}"
    assert os.path.exists(os.path.join(base_dir, row["control_image"])), f"Missing control image: {row['control_image']}"

print("✅ 数据集验证通过！")